import argparse
from os import path as osp

from tools.data_converter import lyft_converter as lyft_converter
from tools.data_converter import nuscenes_converter as nuscenes_converter

def waymo_data_prep(root_path,
                    info_prefix,
                    version,
                    out_dir,
                    workers,
                    max_sweeps=5,
                    debug_mode=False):
    """Prepare the info file for waymo dataset.

    Args:
        root_path (str): Path of dataset root.
        info_prefix (str): The prefix of info filenames.
        out_dir (str): Output directory of the generated info file.
        workers (int): Number of threads to be used.
        max_sweeps (int): Number of input consecutive frames. Default: 5 \
            Here we store pose information of these frames for later use.
    """
    from tools.data_converter import waymo_converter as waymo

    
    load_dir = osp.join(root_path, 'waymo_format')
    
    save_dir = osp.join(out_dir, 'waymo')

    converter = waymo.Waymo4GS(
        load_dir,
        save_dir,
        workers=workers,
        test_mode=False)
    if debug_mode:
        print("debug mode:")
        converter.convert_one(0)
    else:
        converter.convert()
    # Generate waymo infos
    # out_dir = osp.join(out_dir, 'kitti_format')
    # kitti.create_waymo_info_file(out_dir, info_prefix, max_sweeps=max_sweeps)
    # create_groundtruth_database(
    #     'WaymoDataset',
    #     out_dir,
    #     info_prefix,
    #     f'{out_dir}/{info_prefix}_infos_train.pkl',
    #     relative_path=False,
    #     with_mask=False)


parser = argparse.ArgumentParser(description='Data converter arg parser')
parser.add_argument('dataset', metavar='kitti', help='name of the dataset')
parser.add_argument(
    '--root-path',
    type=str,
    default='./data/kitti',
    help='specify the root path of dataset')
parser.add_argument(
    '--version',
    type=str,
    default='v1.0',
    required=False,
    help='specify the dataset version, no need for kitti')
parser.add_argument(
    '--max-sweeps',
    type=int,
    default=10,
    required=False,
    help='specify sweeps of lidar per example')
parser.add_argument(
    '--out-dir',
    type=str,
    default='./data/kitti',
    required='False',
    help='name of info pkl')
parser.add_argument('--extra-tag', type=str, default='kitti')
parser.add_argument('--debug_mode', action='store_true', default=False)
parser.add_argument(
    '--workers', type=int, default=4, help='number of threads to be used')
args = parser.parse_args()

if __name__ == '__main__':
    
    if args.dataset == 'waymo':
        waymo_data_prep(
            root_path=args.root_path,
            info_prefix=args.extra_tag,
            version=args.version,
            out_dir=args.out_dir,
            workers=args.workers,
            max_sweeps=args.max_sweeps,
            debug_mode=args.debug_mode)
    else :
        assert False, "Only Waymo can be handled!"